{% set iz_dev = "" %}
{% if cuda.cudnn.target == "devel" -%}
  {% set iz_dev = "-devel" %}
{% endif %}
ARG IMAGE_NAME
FROM ${IMAGE_NAME}:{{ cuda.version.major }}.{{ cuda.version.minor }}-{{ cuda.cudnn.target }}-{{ cuda.os.distro }}{{ cuda.os.version }}{{ cuda.image_tag_suffix }} as base

{% if "x86_64" in cuda %}
FROM base as base-amd64

{% set cudnn_version = cuda.x86_64.cudnn.version %}
{% if "-" in cuda.x86_64.cudnn.version %}
  {% set cudnn_version = cuda.x86_64.cudnn.version[:-2] %}
{% endif -%}

ENV NV_CUDNN_VERSION {{ cuda.x86_64.cudnn.version }}
ENV NV_CUDNN_PACKAGE libcudnn{{ cuda.x86_64.cudnn.version[0] }}{{ iz_dev }}-${NV_CUDNN_VERSION}.cuda{{ cuda.version.major }}.{{ cuda.version.minor }}

{% if "sha256sum" in cuda.x86_64.cudnn %}
    {% if cuda.version.major_minor == "9.1" %}
        {% set cudnn_url_version = cuda.x86_64.cudnn.version %}
        {% set cudnn_url_version2 = cuda.x86_64.cudnn.version[:cuda.x86_64.cudnn.version.rfind(".")] %}
    {% elif cuda.version.major_minor == "8.0" %}
        {% if cudnn_version[0] | int == 5 or cudnn_version[0] | int == 6 %}
            {% set cudnn_url_version = cuda.x86_64.cudnn.version[:cuda.x86_64.cudnn.version.rfind(".")] %}
            {% set cudnn_url_version2 = cuda.x86_64.cudnn.version[:cuda.x86_64.cudnn.version.rfind(".")] %}
        {% elif cudnn_version[0] | int == 7 %}
            {% set cudnn_url_version = cuda.x86_64.cudnn.version[:cuda.x86_64.cudnn.version.rfind(".")] %}
            {% set cudnn_url_version2 = cuda.x86_64.cudnn.version[:cuda.x86_64.cudnn.version.rfind(".")-2] %}
        {% else %}
            {% set cudnn_url_version = cuda.x86_64.cudnn.version[:3] %}
            {% set cudnn_url_version2 = cuda.x86_64.cudnn.version[:cuda.x86_64.cudnn.version.rfind(".")] %}
        {% endif %}
    {% else %}
        {% set cudnn_url_version = cuda.x86_64.cudnn.version[:cuda.x86_64.cudnn.version.rfind(".")] %}
        {% set cudnn_url_version2 = cudnn_version %}
    {% endif %}
ENV NV_CUDNN_DL_SHA256SUM {{ cuda.x86_64.cudnn.sha256sum }}
ENV NV_CUDNN_DL_BASENAME cudnn-{{ cuda.version.major }}.{{ cuda.version.minor }}-linux-x64-v{{ cudnn_url_version2 }}.tgz
ENV NV_CUDNN_DL_URL http://developer.download.nvidia.com/compute/redist/cudnn/v{{ cudnn_url_version }}/${NV_CUDNN_DL_BASENAME}
{% if cuda.x86_64.cudnn.target == "runtime" -%}ENV NV_CUDNN_DL_RUNTIME_EXTRACT_ARG "--wildcards 'cuda/lib64/libcudnn.so.*'" {% endif -%}
{% elif "sha256sum" not in cuda.x86_64.cudnn %}
    {% set cudnn_package = 1 %}
{% endif %}

{% endif %}
{% if "ppc64le" in cuda %}
FROM base as base-ppc64le

{% set cudnn_version = cuda.ppc64le.cudnn.version %}
{% if "-" in cuda.ppc64le.cudnn.version %}
  {% set cudnn_version = cuda.ppc64le.cudnn.version[:-2] %}
{% endif -%}

ENV NV_CUDNN_VERSION {{ cuda.ppc64le.cudnn.version }}
ENV NV_CUDNN_PACKAGE libcudnn{{ cuda.ppc64le.cudnn.version[0] }}{{ iz_dev }}-${NV_CUDNN_VERSION}.cuda{{ cuda.version.major }}.{{ cuda.version.minor }}

{% if "sha256sum" in cuda.ppc64le.cudnn %}
    {% if cuda.version.major_minor == "9.1" %}
        {% set cudnn_url_version = cuda.ppc64le.cudnn.version %}
        {% set cudnn_url_version2 = cuda.ppc64le.cudnn.version[0] %}
    {% elif cuda.version.major_minor == "8.0" %}
        {% if cudnn_version[0] | int == 5 or cudnn_version[0] | int == 6 %}
            {% set cudnn_url_version = cuda.ppc64le.cudnn.version[:cuda.ppc64le.cudnn.version.rfind(".")] %}
            {% set cudnn_url_version2 = cuda.ppc64le.cudnn.version[:cuda.ppc64le.cudnn.version.rfind(".")] %}
        {% elif cudnn_version[0] | int == 7 %}
            {% set cudnn_url_version = cuda.ppc64le.cudnn.version[:cuda.ppc64le.cudnn.version.rfind(".")] %}
            {% set cudnn_url_version2 = cuda.ppc64le.cudnn.version[:cuda.ppc64le.cudnn.version.rfind(".")-4] %}
        {% else %}
            {% set cudnn_url_version = cuda.ppc64le.cudnn.version[:3] %}
            {% set cudnn_url_version2 = cuda.ppc64le.cudnn.version[0] %}
        {% endif %}
    {% else %}
        {% set cudnn_url_version = cuda.ppc64le.cudnn.version[:cuda.ppc64le.cudnn.version.rfind(".")] %}
        {% set cudnn_url_version2 = cudnn_version %}
    {% endif %}
ENV NV_CUDNN_DL_SHA256SUM {{ cuda.ppc64le.cudnn.sha256sum }}
ENV NV_CUDNN_DL_BASENAME cudnn-{{ cuda.version.major }}.{{ cuda.version.minor }}-linux-ppc64le-v{{ cudnn_url_version2 }}.tgz
ENV NV_CUDNN_DL_URL http://developer.download.nvidia.com/compute/redist/cudnn/v{{ cudnn_url_version }}/${NV_CUDNN_DL_BASENAME}
{% if cuda.ppc64le.cudnn.target == "runtime" -%}ENV NV_CUDNN_DL_RUNTIME_EXTRACT_ARG "--wildcards 'cuda/targets/ppc64le-linux/lib/libcudnn.so.*'"{% endif -%}
{% elif "sha256sum" not in cuda.ppc64le.cudnn %}
    {% set cudnn_package = 1 %}
{% endif %}
{% endif %}

FROM base-${TARGETARCH}

ARG TARGETARCH

LABEL maintainer "NVIDIA CORPORATION <sw-cuda-installer@nvidia.com>"

ENV CUDNN_VERSION ${NV_CUDNN_VERSION}
LABEL com.nvidia.cudnn.version="${CUDNN_VERSION}"

{% if cudnn_package  %}
RUN yum install -y \
    ${NV_CUDNN_PACKAGE} \
    && yum clean all \
    && rm -rf /var/cache/yum/*
{% endif %}

{% if not cudnn_package %}
# cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
RUN curl -fsSL ${NV_CUDNN_DL_URL} -O && \
    echo "${NV_CUDNN_DL_SHA256SUM}  ${NV_CUDNN_DL_BASENAME}" | sha256sum -c - && \
    gunzip ${NV_CUDNN_DL_BASENAME} && \
    tar --no-same-owner -xf ${NV_CUDNN_DL_BASENAME:0:-4}.tar -C /usr/local ${NV_CUDNN_DL_RUNTIME_EXTRACT_ARG} && \
    rm ${NV_CUDNN_DL_BASENAME:0:-4}.tar && \
    ldconfig
{% endif %}
